import logging
import numpy as np

from util.util_log import setup_logging
from FitnessMultiBeamTrace import FitnessMultiBeamTrace


# 配置日志，默认打印到控制台，也可以设置打印到文件
setup_logging()
# setup_logging(log_file="../files/multi-beam-trace/beam2/QE/L=L1+L4/trace_2025-03-19-1644-16beam-qe.txt")
# 获取日志记录器并记录日志
logger = logging.getLogger("[BPSO]")


class BinaryParticle:
    def __init__(self, dimensions):
        self.position = np.random.randint(2, size=dimensions)
        self.velocity = np.zeros(dimensions)
        self.best_position = np.copy(self.position)
        self.best_value = float('inf')

    def flatten(self):
        """Flatten the 2D position array to a 1D array."""
        return self.position.flatten()

    def reshape(self, shape):
        """Reshape the 1D position array back to a 2D array with the given shape."""
        self.position = self.position.reshape(shape)


class BPSO:
    def __init__(self, objective_function, num_particles, dimensions, w=0.7, c1=1.5, c2=1.5):
        """
        初始化BPSO算法。

        :param objective_function: 目标函数
        :param num_particles: 粒子数量
        :param dimensions: 搜索空间维度 (tuple of two integers for 2D arrays)
        :param w: 惯性权重
        :param c1: 认知系数
        :param c2: 社会系数
        """
        self.objective_function = objective_function
        self.num_particles = num_particles
        self.dimensions = dimensions
        self.w = w
        self.c1 = c1
        self.c2 = c2
        self.particles = []
        self.global_best_position = None
        self.global_best_value = float('inf')

    def __generate_population_part(self, before_array, current_array, max_num):
        part = []
        # 添加before和current作为前两个元素
        part.append(before_array.copy().flatten())
        part.append(current_array.copy().flatten())
        # 计算 diff 数组
        diff = before_array != current_array
        # 获取所有 False 的坐标
        false_coords = np.argwhere(diff)
        # 生成剩余的个体
        for _ in range(max_num - 2):
            # 复制 current_array
            modified_array = np.copy(current_array).flatten()
            # 随机选择一些差异行
            if np.sum(diff) > 0:
                # 随机选择一部分 False 坐标
                num_to_select = np.random.randint(0, len(false_coords) + 1)
                selected_false_coords = false_coords[np.random.choice(len(false_coords), num_to_select, replace=False)]
                # 将 modified_array 中选取坐标替换为 before_array
                for coord in selected_false_coords:
                    row, col = coord
                    modified_array[row * self.dimensions[1] + col] = before_array[row, col]
            part.append(modified_array)
        return part

    def initialize_population(self):
        """Initialize the population of particles."""
        self.particles = [BinaryParticle(np.prod(self.dimensions)) for _ in range(self.num_particles)]
        self.global_best_position = None
        self.global_best_value = float('inf')

    def initialize_population_group(self, phaseBit_mix_ps_before, phaseBit_mix_ps, phaseBit_mix_qe_before,
                                    phaseBit_mix_qe):
        parts_ps = self.__generate_population_part(phaseBit_mix_ps_before, phaseBit_mix_ps, self.num_particles // 2)
        parts_qe = self.__generate_population_part(phaseBit_mix_qe_before, phaseBit_mix_qe,
                                                 self.num_particles - self.num_particles // 2)
        self.particles = [BinaryParticle(np.prod(self.dimensions)) for _ in range(self.num_particles)]
        for i, pos in enumerate(parts_ps + parts_qe):
            self.particles[i].position = pos
        self.global_best_position = None
        self.global_best_value = float('inf')

    def optimize(self, max_iterations, phaseBit_mix_ps_before, phaseBit_mix_ps, phaseBit_mix_qe_before, phaseBit_mix_qe):
        """
        运行BPSO算法进行优化。

        :param max_iterations: 最大迭代次数
        :param phaseBit_mix_ps_before: 上一个时刻的相位比特矩阵 (ps)
        :param phaseBit_mix_ps: 当前时刻的相位比特矩阵 (ps)
        :param phaseBit_mix_qe_before: 上一个时刻的相位比特矩阵 (qe)
        :param phaseBit_mix_qe: 当前时刻的相位比特矩阵 (qe)
        """
        # """初始化返回值"""
        self.best_individual = phaseBit_mix_ps
        self.best_fitness = self.objective_function(self.best_individual)
        self.best_individual_history = []
        self.best_fitness_history = []
        self.global_best_position = self.best_individual
        self.global_best_value = self.best_fitness

        # 初始化
        self.initialize_population_group(phaseBit_mix_ps_before, phaseBit_mix_ps,
                                         phaseBit_mix_qe_before, phaseBit_mix_qe)

        for iteration in range(max_iterations):
            for particle in self.particles:
                # Reshape the flattened position back to 2D array
                particle.position = particle.position.reshape(self.dimensions)

                # Calculate the current value of the objective function
                current_value = self.objective_function(particle.position)

                # Update the best position and value for the particle
                if current_value < particle.best_value:
                    particle.best_value = current_value
                    particle.best_position = np.copy(particle.position)

                    # Update the global best position and value if needed
                    if current_value < self.global_best_value:
                        self.global_best_value = current_value
                        self.global_best_position = np.copy(particle.position)

            # Flatten the positions for velocity updates
            for particle in self.particles:
                particle.position = particle.flatten()

            # Update velocities and positions
            for particle in self.particles:
                r1 = np.random.rand(len(particle.position))
                r2 = np.random.rand(len(particle.position))

                cognitive_velocity = self.c1 * r1 * (particle.best_position.flatten() - particle.position)
                social_velocity = self.c2 * r2 * (self.global_best_position.flatten() - particle.position)

                particle.velocity = self.w * particle.velocity + cognitive_velocity + social_velocity

                # Update particle position using sigmoid transfer function and binarization
                sigmoid_prob = 1 / (1 + np.exp(-particle.velocity))
                particle.position = (np.random.rand(len(particle.position)) < sigmoid_prob).astype(int)

            # 更新最佳适应度
            if self.global_best_value < self.best_fitness:
                self.best_fitness = self.global_best_value
                self.best_individual = self.global_best_position
            # 记录最佳适应度曲线
            self.best_fitness_history.append(self.best_fitness)
            self.best_individual_history.append(self.best_individual)

            # Print the best solution found so far
            # print(
            #     f"Iteration {iteration + 1}: Best position = {self.global_best_position}, Best value = {self.global_best_value}")
            logger.info("iteration=%d: self.best_fitness=%f" % (iteration, self.best_fitness))
            logger.info("---------------------------------------------------------------------------")

        # return self.global_best_position, self.global_best_value
        return self.best_individual, self.best_fitness, self.best_fitness_history, self.best_individual_history




def test():
    # 使用BPSO算法优化最大化一的数量的问题
    # fixme: 需要修改util_ris_pattern.py
    # # 初始化参数
    # nRow = 64  # x方向单元个数，θ方向
    # mLine = 64  # y方向单元个数，φ方向
    dimensions = (3, 3)  # Example 3x3 matrix
    num_particles = 30
    max_iterations = 10

    # 创建一个示例 FitnessMultiBeamTrace 实例
    bit_num = 1  # 根据实际情况设置比特数
    beam_num = 1  # 根据实际情况设置波束数
    fitness_function = FitnessMultiBeamTrace(bit_num, beam_num)

    # 初始化主瓣位置
    phaseBit_target = np.random.randint(2, size=(3, 3))
    print(f"phaseBit_target: {phaseBit_target}")
    fitness_function.init_main_lope_pos(phaseBit_target)

    phaseBit_ps_before = np.array([[0, 0, 0], [0, 0, 0], [0, 0, 0]])
    phaseBit_ps = np.array([[0, 0, 0], [0, 0, 0], [1, 1, 1]])
    phaseBit_qe_before = np.array([[0, 0, 0], [0, 0, 0], [0, 0, 0]])
    phaseBit_qe = np.array([[1, 1, 1], [0, 0, 0], [0, 0, 0]])

    bpso = BPSO(objective_function=fitness_function.fitness, num_particles=num_particles, dimensions=dimensions)
    best_position, best_value, best_position_list, best_value_list \
        = bpso.optimize(max_iterations, phaseBit_ps_before, phaseBit_ps, phaseBit_qe_before, phaseBit_qe)
    print(f"\nFinal Best position:\n{best_position}")
    print(f"Final Best value: {-best_value}")  # Convert back to maximizing the sum of ones
    print("--------------------------------------------------------------------------------------")
    best_position, best_value, best_position_list, best_value_list \
        = bpso.optimize(max_iterations, phaseBit_ps_before, phaseBit_ps, phaseBit_qe_before, phaseBit_qe)
    print(f"\nFinal Best position:\n{best_position}")
    print(f"Final Best value: {-best_value}")  # Convert back to maximizing the sum of ones



if __name__=="__main__":
    test()